package com.lark.algorithm.study.trie;

import java.util.TreeMap;

/**
 * @author btmood
 * @version 1.0
 * @apiNote 字典树
 * @since 2024-06-14 14:06
 */
public class Trie {

    private class Node {
        // 一个单词是否结束
        public boolean isWord;
        /**
         * 子节点的映射
         * 这里把字符设置成了Character，是因为英文字母是字符
         * 但是诸如汉语等语言，单个字就不是字符了，可以换成其它类型
         * 也就是说只要能拆分的文本结构，类型不同，都可以使用Trie来构建
         * 后期可以改造成泛型
         */
        public TreeMap<Character, Node> next;

        public Node(boolean isWord) {
            this.isWord = isWord;
            next = new TreeMap<>();
        }

        public Node() {
            this(false);
        }
    }

    // 根节点
    private Node root;
    // 容量
    private int size;

    public Trie() {
        root = new Node();
        size = 0;
    }

    /**
     * 获取容量
     * @return
     */
    public int getSize() {
        return size;
    }

    /**
     * 向Trie中添加一个新的单词word
     * @param word 单词
     */
    public void add(String word) {
        Node cur = root;
        for (char c : word.toCharArray()) {
            if (cur.next.get(c) == null) {
                cur.next.put(c, new Node());
            }
            cur = cur.next.get(c);
        }
        if (!cur.isWord) {
            cur.isWord = true;
            size++;
        }
    }

    /**
     * 向Trie中添加一个新的单词word（递归写法）
     * @param word 单词
     */
    public void addRecursion(String word) {
        doAddRecursion(root, word);
    }

    /**
     * 向以node为根的Trie中添加单词word
     * @param node 根
     * @param word 单词
     */
    public void doAddRecursion(Node node, String word) {
        if (word.isEmpty()) {
            // 别忘了把最后一个字符设置标记
            if (!node.isWord) {
                node.isWord = true;
                size++;
            }
            return;
        }
        Character c = word.charAt(0);
        if (node.next.get(c) == null) {
            node.next.put(c, new Node());
        }
        node = node.next.get(c);
        doAddRecursion(node, word.substring(1));
    }

    /**
     * 查询单词word是否在Trie中（非递归写法）
     * @param word 单词
     * @return 布尔
     */
    public boolean contains(String word) {
        Node cur = root;
        for (Character c : word.toCharArray()) {
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }
        // 这里要注意，如果最后一个字符不是word，说明这个单词并没有添加到Trie中，只能返回false
        return cur.isWord;
    }

    /**
     * 查询单词word是否在Trie中（递归写法）
     * @param word 单词
     * @return 布尔
     */
    public boolean containsRecursion(String word) {
        return doContainsRecursion(root, word);
    }

    /**
     * 查询以node为根的Trie中是否存在单词word
     * @param node 根
     * @param word 单词
     * @return 布尔
     */
    public boolean doContainsRecursion(Node node, String word) {
        if (word.isEmpty()) {
            return node.isWord;
        }
        char c = word.charAt(0);
        if (node.next.get(c) == null) {
            return false;
        }
        return doContainsRecursion(node.next.get(c), word.substring(1));
    }

    /**
     * 查询是否在Trie中有单词以prefix为前缀
     * 一个单词也是这个单词本身的前缀
     * @param prefix 前缀
     * @return 布尔
     */
    public boolean isPrefix(String prefix) {
        Node cur = root;
        for (Character c : prefix.toCharArray()) {
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }
        return true;
    }



}
